import sys
from PyQt5 import QtWidgets
import paramiko
import io
import os
from scp import SCPClient

class StreamToTextEdit(io.StringIO):
    def __init__(self, text_edit):
        super().__init__()
        self.text_edit = text_edit

    def write(self, message):
        self.text_edit.append(message)

    def flush(self):
        pass  # 不需要实现

class FileListWindow(QtWidgets.QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("文件选择")
        self.setGeometry(100, 100, 600, 400)

        self.local_directory = ""  # 初始化本地目录
        self.UbuntuIp = "192.168.3.103"  # 默认 IP 地址
        self.SetUp()
    
    def UIinit(self):
        # 创建主水平布局
        self.main_layout = QtWidgets.QHBoxLayout()

        # 左边区域
        self.left_layout = QtWidgets.QVBoxLayout()
        
        self.button_select_path = QtWidgets.QPushButton("选择路径")
        self.button_update_list = QtWidgets.QPushButton("更新列表")
        self.button_download_file = QtWidgets.QPushButton("下载文件")
        self.button_set_ip = QtWidgets.QPushButton("设置 IP 地址")  # 新增按钮
        
        self.left_layout.addWidget(self.button_select_path)
        self.left_layout.addWidget(self.button_update_list)
        self.left_layout.addWidget(self.button_download_file)
        self.left_layout.addWidget(self.button_set_ip)  # 添加到布局

        # 文本输出框
        self.output_text = QtWidgets.QTextEdit()
        self.output_text.setReadOnly(True)  # 设置为不可编辑
        self.left_layout.addWidget(self.output_text)

        # 右边区域
        self.right_layout = QtWidgets.QVBoxLayout()
        self.scroll_area = QtWidgets.QScrollArea(self)
        self.scroll_area.setWidgetResizable(True)

        # 内容小部件
        self.content_widget = QtWidgets.QWidget()
        self.content_layout = QtWidgets.QVBoxLayout()
        self.content_widget.setLayout(self.content_layout)
        self.scroll_area.setWidget(self.content_widget)

        self.right_layout.addWidget(self.scroll_area)

        # 将左右布局添加到主布局
        self.main_layout.addLayout(self.left_layout, 1)
        self.main_layout.addLayout(self.right_layout, 3)

        self.setLayout(self.main_layout)

    def UIqss(self):
        # 重定向输出
        sys.stdout = StreamToTextEdit(self.output_text)

    def UIfun(self):
        self.button_select_path.clicked.connect(self.select_path)
        self.button_update_list.clicked.connect(self.UpdateList)
        self.button_download_file.clicked.connect(self.download_files)  # 连接下载按钮
        self.button_set_ip.clicked.connect(self.set_ip_address)  # 连接设置 IP 按钮

    def get_file_list(self, ssh, remote_directory):
        stdin, stdout, stderr = ssh.exec_command(f'ls {remote_directory}')
        return stdout.read().decode().splitlines()

    def UpdateList(self):
        try:
            # SSH 连接信息
            hostname = self.UbuntuIp
            username = "sts"  # 用户名
            password = "5210"  # 密码
            remote_directory = "~/Desktop/vex-timer/savedata"  # 远程文件夹路径

            # 创建 SSH 客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname, username=username, password=password)

            # 获取文件列表
            file_list = self.get_file_list(ssh, remote_directory)

            # 清空之前的复选框
            for i in reversed(range(self.content_layout.count())): 
                widget = self.content_layout.itemAt(i).widget()
                if widget is not None:
                    widget.deleteLater()

            for file in file_list:
                checkbox = QtWidgets.QCheckBox(file)
                checkbox.setProperty("remote_file", file)  # 保存文件名
                self.content_layout.addWidget(checkbox)
        except Exception as e:
            print(f"出错了，检查 IP 是否正确: {e}")

    def select_path(self):
        # 打开文件夹选择对话框
        folder = QtWidgets.QFileDialog.getExistingDirectory(self, "选择保存路径")
        if folder:
            self.local_directory = folder
            print(f"选择的路径: {self.local_directory}")

    def set_ip_address(self):
        # 弹出对话框输入新 IP
        text, ok = QtWidgets.QInputDialog.getText(self, "设置 IP 地址", "输入新的 Ubuntu IP 地址:", QtWidgets.QLineEdit.Normal, self.UbuntuIp)
        if ok and text:
            self.UbuntuIp = text
            print(f"新的 IP 地址: {self.UbuntuIp}")

    def download_files(self):
        if not self.local_directory:
            print("请先选择保存路径.")
            return
        
        try:
            hostname = self.UbuntuIp
            username = "sts"  # 用户名
            password = "5210"  # 密码

            # 创建 SSH 客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname, username=username, password=password)
            remote_directory = "~/Desktop/vex-timer/savedata/"  # 远程文件夹路径

            # 创建 SCP 客户端
            with SCPClient(ssh.get_transport()) as scp:
                # 下载已勾选的文件
                for i in range(self.content_layout.count()):
                    checkbox = self.content_layout.itemAt(i).widget()
                    if isinstance(checkbox, QtWidgets.QCheckBox) and checkbox.isChecked():
                        remote_file = checkbox.property("remote_file")  # 获取文件名
                        local_file_path = os.path.join(self.local_directory, remote_file)
                        print(f"下载文件: {remote_file} 到 {local_file_path}")
                        scp.get(remote_directory + remote_file, local_file_path)  # 下载文件

            print("下载完成!")
        except Exception as e:
            print(f"下载过程中出错: {e}")

    def SetUp(self):
        self.UIinit()
        self.UIqss()
        self.UIfun()

def main():
    app = QtWidgets.QApplication(sys.argv)
    window = FileListWindow()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()